#include "elfloader.h"
#include "elf64.h"
#include "bit.h"

static mx_usize_t dyn_memsize (mx_elf64_analysis_pt panalysis)
{
    mx_usize_t ret = 0;
#ifdef MXSETUP_KAPP_INTEGRATE_SYM
    /* the memory size = __bss_end - baseaddr */
    mx_uint64_t baseaddr = panalysis->shdr.start[1].sh_addr;
    mx_elf64_sym_pt sym_bss_end = mx_elf64_symbol_getbyname(panalysis, "__kbss_end");
    if (sym_bss_end != mx_null)
    {
        mx_uint64_t bss_end = mx_elf64_symbol_vaddr(panalysis, sym_bss_end);
        ret = bss_end - baseaddr;
    }
#else
    #error "not support"
#endif /* MXSETUP_KAPP_INTEGRATE_SYM */
    return ret;
}

/**
 * calcute the memory size of elf that it need
 * @param elf: the elf binary
 * @return: memory size
 */
mx_usize_t mx_elfloader_memsize (mx_uint8_pt elf)
{
    mx_usize_t ret = 0;
    mx_elf64_analysis_t analysis;
    if (!mx_elf64_analysis_init(&analysis, elf))
    {
        return ret;
    }
    /* calculat the execution memory size under different ELF type */
    switch (analysis.ehdr->e_type)
    {
    case ELF64_EHDR_E_TYPE_REL:
        break;
    case ELF64_EHDR_E_TYPE_EXEC:
        break;
    case ELF64_EHDR_E_TYPE_DYN:
        ret = dyn_memsize(&analysis);
        break;
    }
    return ret;
}

/* load and relocation repaire of Dynamic Object (DYN)
   return the base address (MX_ELFLOADER_BASEADDR) and 0 mean that failed */
static mx_uint64_t dyn_load8relocation (mx_elf64_analysis_pt panalysis, mx_uint64_t addr)
{
    mx_uint64_t baseaddr;
#ifdef MXSETUP_KAPP_INTEGRATE_SYM
    /* get the base address (MX_ELFLOADER_BASEADDR), the execution address of shdr[1] is minimum
       and also equal MX_ELFLOADER_BASEADDR.
       NOTE: the MX_ELFLOADER_BASEADDR can't be used directly because KAPP may be compiled with different kconfig */
    baseaddr = panalysis->shdr.start[1].sh_addr;
    /* only load sections which sh_flags == PROGBITS,
       and NOBITS or other section who generate by ld needn't load */
    mx_elf64_shdr_foreach(panalysis, shdr, {
        if (GETMSK(shdr->sh_flags, ELF64_SHDR_SH_FLAGS_ALLOC_MASK) &&
            shdr->sh_type == ELF64_SHDR_SH_TYPE_PROGBITS)
        {
            mx_uint64_t offset = shdr->sh_addr - baseaddr;
            // FIXME: copying by bytes is too slow
            mx_memcpy((mx_uint8_pt)((mx_uint64_t)addr + offset),
                      (mx_uint8_pt)mx_elf64_section_paddr(panalysis, shdr), shdr->sh_size);
        }
    });
    /* do relocation repaire */
    mx_elf64_shdr_foreach(panalysis, shdr, {
        if (shdr->sh_type == ELF64_SHDR_SH_TYPE_RELA)
        {
            mx_elf64_rela_foreach(panalysis, shdr, rela, {
                /* get the virtual address of the symbol corresponding to the rela */
                mx_uint64_t symaddr = mx_elf64_rela_symvaddr(panalysis, shdr, rela);
                /* get the virtual address which need to repair */
                mx_uint64_t repaddr = mx_elf64_relx_repvaddr(panalysis, shdr, rela);
                /* obtain the symbol table associated with this relocation entry */
                mx_elf64_sym_pt sym = mx_elf64_relx_symbol(panalysis, shdr, rela);
                /* if option `MXSETUP_KAPP_INTEGRATE_SYM` is used, the existence of undifined symbol is not allowed */
                if (sym->st_shndx == ELF64_SYM_ST_SHNDX_UNDEF)
                {
                    /* the synbol dosen't have a corresponging section, means that it is a undefined symbol,
                       it is an illegal situation, so that exit with mx_null */
                    baseaddr = 0;
                    break;
                }
                /* if a symbol's type is ABS means that it is a kernel symbol, so ABS symbol needn't
                   calculate the new address after reload, otherwise we need calculate the new execution address */
                if (sym->st_shndx != ELF64_SYM_ST_SHNDX_ABS)
                {
                    /* new address + offset */
                    symaddr = addr + (symaddr - baseaddr);
                }
                /* calculate the new address where need to repair after load. new address + offset */
                repaddr = addr + (repaddr - baseaddr);
                /* do repair (architecture-related) */
                if (!mx_elfloader_relx_repair(rela->r_info.type, repaddr, symaddr))
                {
                    /* the corresponding architecture doesn't support this type of relocation */
                    baseaddr = 0;
                    break;
                }
            })
        }
        else if (shdr->sh_type == ELF64_SHDR_SH_TYPE_REL)
        {
            /* TODO: to be supplemented */
            baseaddr = 0;
            break;
        }
    });
#else
    #error "not support"
#endif /* MXSETUP_KAPP_INTEGRATE_SYM */
    return baseaddr;
}

/**
 * load elf to the specified address
 * @param elf: the elf binary
 * @param addr: the address that the elf load to
 * @return: the entry point of the elf
 */
kmain_pt mx_elfloader (mx_uint8_pt elf, mx_void_pt addr)
{
    mx_elf64_analysis_t analysis;
    if (!mx_elf64_analysis_init(&analysis, elf))
    {
        return mx_null;
    }
    /* check the `.kapp.label` section */
    mx_elf64_shdr_pt shdr_label = mx_null;
    mx_elf64_shdr_foreach(&analysis, shdr, {
        if (mx_strcmp(mx_elf64_shdr_name(&analysis, shdr), ".kapp.label") == 0)
        {
            shdr_label = shdr;
            break;
        }
    });
    if (shdr_label == mx_null)
    {
        /* section `.kapp.label` not found */
        return mx_null;
    }
    /* check the magic */
    mx_elfloader_label_pt label = (mx_elfloader_label_pt)mx_elf64_section_paddr(&analysis, shdr_label);
    if (label->magic != MX_ELFLOADER_LABEL_MAGIC)
    {
        /* error magic value */
        return mx_null;
    }
    mx_uint64_t baseaddr = 0;
    /* load elf to memory do relocation under different ELF type */
    switch (analysis.ehdr->e_type)
    {
    case ELF64_EHDR_E_TYPE_REL:
        break;
    case ELF64_EHDR_E_TYPE_EXEC:
        break;
    case ELF64_EHDR_E_TYPE_DYN:
        baseaddr = dyn_load8relocation(&analysis, (mx_uint64_t)addr);
        break;
    }
    kmain_pt kmain = mx_null;
    /* if relocation repair successed, then do .bss clean operation
       and get the entry point (main) of kapp if operation was successful */
    if (baseaddr != 0)
    {
        /* do .bss clean operation, get the region of .bss by relevant symbols in .lds */
        mx_elf64_sym_pt sym_bss_start = mx_elf64_symbol_getbyname(&analysis, "__kbss_start");
        mx_elf64_sym_pt sym_bss_end = mx_elf64_symbol_getbyname(&analysis, "__kbss_end");
        if (sym_bss_start != mx_null && sym_bss_end != mx_null)
        {
            /* get the execution address of symbols and reposition their */
            mx_uint64_t bss_start = mx_elf64_symbol_vaddr(&analysis, sym_bss_start);
            mx_uint64_t bss_end = mx_elf64_symbol_vaddr(&analysis, sym_bss_end);
            bss_start = (mx_uint64_t)addr + (bss_start - baseaddr);     // new address + offset
            bss_end = (mx_uint64_t)addr + (bss_end - baseaddr);
            /* fill 0 into .bss region */
            mx_memset((mx_uint8_pt)bss_start, 0, bss_end - bss_start);  // FIXME: copying by bytes is too slow
            /* get the entry point of kapp */
            kmain = (kmain_pt)((mx_uint64_t)addr + (analysis.ehdr->e_entry - baseaddr));
        }
    }
    return kmain;
}
